import streamlit as st
import logging
from typing import Any, Dict, Optional, List
from datetime import datetime
from pathlib import Path
import json

class StateManager:
    """状态管理器"""
    def __init__(self, session_dir: str = "chat_sessions"):
        self.session_dir = Path(session_dir)
        self.session_dir.mkdir(parents=True, exist_ok=True)
        self.logger = logging.getLogger(__name__)
        self._initialize_states()
    
    def _initialize_states(self):
        """初始化所有状态"""
        # 基础状态
        if "messages" not in st.session_state:
            st.session_state.messages = []
        if "model" not in st.session_state:
            st.session_state.model = "llama2"
        if "temperature" not in st.session_state:
            st.session_state.temperature = 0.7
        if "max_length" not in st.session_state:
            st.session_state.max_length = 2000
        if "current_session" not in st.session_state:
            st.session_state.current_session = None
            
        # 会话管理
        if "sessions" not in st.session_state:
            st.session_state.sessions = self._load_sessions()
        if "show_delete_confirm" not in st.session_state:
            st.session_state.show_delete_confirm = None
        if "show_menu" not in st.session_state:
            st.session_state.show_menu = None
            
        # 服务器配置
        if "server_ip" not in st.session_state:
            st.session_state.server_ip = "localhost"
        if "server_port" not in st.session_state:
            st.session_state.server_port = "11434"
        if "found_servers" not in st.session_state:
            st.session_state.found_servers = []
            
        # 网络扫描
        if "custom_networks" not in st.session_state:
            st.session_state.custom_networks = set()
            
        # UI状态
        if "current_response" not in st.session_state:
            st.session_state.current_response = None
        if "show_settings" not in st.session_state:
            st.session_state.show_settings = False
        if "show_history" not in st.session_state:
            st.session_state.show_history = False
    
    def _load_sessions(self) -> Dict:
        """加载所有会话"""
        sessions = {}
        try:
            for file in self.session_dir.glob("*.json"):
                session_name = file.stem
                with open(file, 'r', encoding='utf-8') as f:
                    sessions[session_name] = json.load(f)
        except Exception as e:
            self.logger.error(f"Failed to load sessions: {str(e)}")
        return sessions
    
    def save_session(self, session_name: str):
        """保存指定会话"""
        try:
            if session_name and session_name in st.session_state.sessions:
                file_path = self.session_dir / f"{session_name}.json"
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(st.session_state.sessions[session_name], f, 
                            ensure_ascii=False, indent=2)
        except Exception as e:
            self.logger.error(f"Failed to save session {session_name}: {str(e)}")
    
    def create_session(self, session_name: str = None):
        """创建新会话"""
        if not session_name:
            session_name = f"会话 {datetime.now().strftime('%Y%m%d_%H%M%S')}"
            
        sessions = self.get_state("sessions", {})
        sessions[session_name] = {
            "created_at": datetime.now().isoformat(),
            "messages": []
        }
        
        self.set_state("sessions", sessions)
        self.set_state("current_session", session_name)
        self.set_state("messages", [])
    
    def switch_session(self, session_name: str) -> bool:
        """切换会话"""
        try:
            if session_name != self.get_state("current_session"):
                # 保存当前会话
                self.save_current_session()
                
                # 切换会话
                st.session_state.current_session = session_name
                st.session_state.messages = st.session_state.sessions[session_name].copy()
                return True
        except Exception as e:
            self.logger.error(f"Failed to switch session: {str(e)}")
        return False
    
    def delete_session(self, session_id: str) -> bool:
        """删除会话"""
        try:
            if session_id == self.get_state("current_session"):
                self.update_states({
                    "current_session": None,
                    "messages": []
                })
            
            session_file = self.session_dir / f"{session_id}.json"
            if session_file.exists():
                session_file.unlink()
            return True
        except Exception as e:
            self.logger.error(f"删除会话失败: {str(e)}")
            return False
    
    def get_state(self, key: str, default: Any = None) -> Any:
        """获取状态值"""
        return getattr(st.session_state, key, default)
    
    def set_state(self, key: str, value: Any):
        """设置状态值"""
        setattr(st.session_state, key, value)
    
    def update_states(self, states: Dict[str, Any]):
        """批量更新状态
        
        Args:
            states: 要更新的状态字典
        """
        for key, value in states.items():
            self.set_state(key, value)
    
    def save_current_session(self):
        """保存当前会话"""
        current_session = self.get_state("current_session")
        if current_session:
            st.session_state.sessions[current_session] = st.session_state.messages.copy()
            self.save_session(current_session)
    
    def clear_current_session(self):
        """清空当前会话"""
        st.session_state.messages = []
        self.save_current_session()
    
    def get_session_list(self) -> List[str]:
        """获取会话列表"""
        return sorted(st.session_state.sessions.keys())
    
    def get_session_info(self, session_name: str) -> Optional[Dict]:
        """获取会话信息"""
        try:
            if session_name in st.session_state.sessions:
                messages = st.session_state.sessions[session_name]
                return {
                    "name": session_name,
                    "message_count": len(messages),
                    "last_message": messages[-1]["content"] if messages else None,
                    "is_current": session_name == self.get_state("current_session")
                }
        except Exception as e:
            self.logger.error(f"Failed to get session info: {str(e)}")
        return None

    def new_session(self) -> str:
        """创建新会话"""
        session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.update_states({
            "current_session": session_id,
            "messages": []
        })
        # 确保清空消息历史
        st.session_state.messages = []
        self._save_session()
        self.logger.info(f"已创建新会话: {session_id}")
        return session_id
    
    def load_session(self, session_id: str) -> bool:
        """加载会话"""
        try:
            session_file = self.session_dir / f"{session_id}.json"
            if session_file.exists():
                data = json.loads(session_file.read_text())
                
                # 更新会话状态
                self.update_states({
                    "messages": data.get("messages", []),
                    "model": data.get("model", "llama2"),
                    "temperature": data.get("temperature", 0.7),
                    "max_length": data.get("max_length", 2000),
                    "current_session": session_id
                })
                
                # 同步到 session_state
                st.session_state.messages = data.get("messages", []).copy()
                
                self.logger.info(f"已加载会话: {session_id}")
                return True
                
        except Exception as e:
            self.logger.error(f"加载会话失败: {str(e)}")
        return False
    
    def _save_session(self):
        """保存当前会话"""
        if self.get_state("current_session"):
            session_file = self.session_dir / f"{self.get_state('current_session')}.json"
            try:
                data = {
                    "messages": self.get_state("messages"),
                    "model": self.get_state("model"),
                    "temperature": self.get_state("temperature"),
                    "max_length": self.get_state("max_length"),
                    "updated_at": datetime.now().isoformat()
                }
                session_file.write_text(json.dumps(data, indent=2))
            except Exception as e:
                self.logger.error(f"保存会话失败: {str(e)}")
    
    def append_message(self, role: str, content: str):
        """添加消息"""
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }
        self.set_state("messages", self.get_state("messages") + [message])
        self._save_session()
    
    def get_messages(self) -> List[Dict]:
        """获取所有消息"""
        return self.get_state("messages")
    
    def clear_messages(self):
        """清空消息"""
        self.set_state("messages", [])
        st.session_state.messages = []
        self._save_session()
        self.logger.info("已清空当前会话")
    
    def get_sessions(self) -> List[str]:
        """获取所有会话"""
        return [
            f.stem for f in self.session_dir.glob("*.json")
        ]
    
    def get_current_session(self) -> Optional[str]:
        """获取当前会话ID"""
        return self.get_state("current_session")
    
    def update_settings(self, settings: Dict):
        """更新设置"""
        self.set_state("model", settings.get("model", "llama2"))
        self.set_state("temperature", settings.get("temperature", 0.7))
        self.set_state("max_length", settings.get("max_length", 2000))
        self._save_session()
    
    def get_settings(self) -> Dict:
        """获取设置"""
        return {
            "model": self.get_state("model"),
            "temperature": self.get_state("temperature"),
            "max_length": self.get_state("max_length")
        } 